# 为了CycleFusion与其他的比较
# 重复之前D2BEGAN的过程

import os
import shutil

# 记录有多少是D（白天）
cnt = 0
# 路径
ir_dir = r"D:\cycleFusion-project\dataset-1444\msrs_c1\ir"
vi_dir = r"D:\cycleFusion-project\dataset-1444\msrs_c1\vi"
target_ir_dir = os.path.join(ir_dir.rsplit('\\', maxsplit=1)[0], "new_dataset/ir")
os.makedirs(target_ir_dir, exist_ok=True)
target_vi_dir = os.path.join(ir_dir.rsplit('\\', maxsplit=1)[0], "new_dataset/vi")
os.makedirs(target_vi_dir, exist_ok=True)

# 创建一个字典，映射每个数字后面是D还是N
name2DN = {n[:-1]:n[-1] for n in map(lambda x:x.split('.')[0],os.listdir(ir_dir))}

def extract(imgsNo):

    cnt = 0

    "抽取字符串imgsNo中的图片，如果是个范围就不要加字母了，如27-38，不能27D-38D"
    for name in imgsNo.split():  # 将imgsNo分割成单个的文件名，split()不加空格表示用空白分割
        name.strip()
        if "-" in name:
            nm = list(name.split('-'))
            for i in range(len(nm) - 1):
                for j in range(eval(nm[i]),eval(nm[i+1]) + 1):
                    cnt += extract_pics(str(j))
        else:
            cnt += extract_pics(name)
    print(f"\nD有{cnt}个")
def extract_pics(name):

    # 补全图片名长度
    if name[-1] in "DN":
        name = '0'*(6-len(name)) + name
    else:
        name = '0'*(5-len(name)) + name
        if name2DN.get(name,None):  # 有的图片名是不存在的。
            name += name2DN[name]
        else:
            return 0

    # cnt += name[-1] == 'D'   未解析引用？？？
    print(name, end=',')
    # 注意，这里是复制
    shutil.copyfile(os.path.join(ir_dir,name + ".png"), os.path.join(target_ir_dir, name + ".png"))  # 复制
    shutil.copyfile(os.path.join(vi_dir,name + ".png"), os.path.join(target_vi_dir, name + ".png"))  # 复制

    return name[-1] == 'D'
if __name__ == '__main__':
    imgsNo = "00003N 00006N 00011N 00012N 00013N 00028N 00031N 00032N 00036N 44N 87D 37N 38N 95D 106D 107D 123D 159-166  170-172-192   202-208  257D 258D  275-277 279  281-285-307  329-338 374D 401D 466D  484D 490D  822N  913 937 940 942-944 954  971  1008  1353 1500 134D  211D 372D  492D 523D 577D  592D  888 909"
    print("开始处理",end=':')
    extract(imgsNo)
